package com.linc.concurrency.effectiveDemo;

import java.util.concurrent.*;

public class Memoizer<A, V> implements Computable<A, V> {
  private final ConcurrentMap<A, Future<V>> cache = new ConcurrentHashMap<>();
  private final Computable<A, V> c;

  public Memoizer(Computable<A, V> c) {
    this.c = c;
  }

  @Override
  public V compute(A args) throws InterruptedException {
    while (true) {
      Future<V> f = cache.get(args);
      if (f == null) {
        Callable<V> eval = () -> c.compute(args);
        FutureTask<V> ft = new FutureTask<>(eval);
        f = cache.putIfAbsent(args, ft);
        if (f == null) {
          f = ft;
          ft.run();
        }
        try {
          return f.get();
        } catch (CancellationException e) {
          cache.remove(args, f);
        } catch (ExecutionException e) {
          var cause = e.getCause();
          if (cause instanceof Error) throw (Error) cause;
          else if (cause instanceof RuntimeException) throw (RuntimeException) cause;
          else throw new IllegalStateException("Not checked", cause);
        }
      }
    }
  }
}
